package com.jvmmonitor.agent;

import java.lang.instrument.Instrumentation;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.jvmmonitor.jmx.JmxMonitor;
import com.jvmmonitor.service.MemoryAnalyzer;
import com.jvmmonitor.service.ObjectSizeMonitor;
import com.jvmmonitor.util.ConfigLoader;
import com.jvmmonitor.util.ObjectSizeCalculator;

/**
 * JVM监控Agent入口类
 */
public class JvmMonitorAgent {
    private static final Logger logger = LoggerFactory.getLogger(JvmMonitorAgent.class);
    
    // 单例实例
    private static JvmMonitorAgent instance;
    
    // 定时执行器
    private ScheduledExecutorService scheduler;
    
    // 监控组件
    private JmxMonitor jmxMonitor;
    private ObjectSizeMonitor objectSizeMonitor;
    private MemoryAnalyzer memoryAnalyzer;
    
    /**
     * Java Agent启动方法（JVM启动时加载）
     * @param args Agent参数
     * @param inst Instrumentation实例
     */
    public static void premain(String args, Instrumentation inst) {
        logger.info("JVM Monitor Agent starting via premain...");
        initialize(inst);
    }
    
    /**
     * Java Agent动态Attach方法（运行时加载）
     * @param args Agent参数
     * @param inst Instrumentation实例
     */
    public static void agentmain(String args, Instrumentation inst) {
        logger.info("JVM Monitor Agent starting via agentmain...");
        initialize(inst);
    }
    
    /**
     * 初始化Agent
     * @param inst Instrumentation实例
     */
    private static synchronized void initialize(Instrumentation inst) {
        if (instance == null) {
            instance = new JvmMonitorAgent();
            instance.init(inst);
        } else {
            logger.warn("JVM Monitor Agent is already initialized");
        }
    }
    
    /**
     * 初始化监控组件
     * @param inst Instrumentation实例
     */
    private void init(Instrumentation inst) {
        try {
            // 设置Instrumentation实例用于对象大小计算
            ObjectSizeCalculator.setInstrumentation(inst);
            
            // 初始化监控组件
            jmxMonitor = new JmxMonitor();
            objectSizeMonitor = new ObjectSizeMonitor();
            memoryAnalyzer = new MemoryAnalyzer();
            
            // 将ObjectSizeMonitor关联到JmxMonitor，用于GC回收对象大小统计
            jmxMonitor.setObjectSizeMonitor(objectSizeMonitor);
            
            // 启动JMX监控
            jmxMonitor.start();
            
            // 设置字节码增强
            setupBytecodeEnhancement(inst);
            
            // 启动定期统计任务
            startScheduledTasks();
            
            logger.info("JVM Monitor Agent initialized successfully");
        } catch (Exception e) {
            logger.error("Failed to initialize JVM Monitor Agent: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 设置字节码增强，监控对象创建
     * @param inst Instrumentation实例
     */
    private void setupBytecodeEnhancement(Instrumentation inst) {
        try {
            logger.info("Setting up bytecode enhancement...");
            
            // 获取需要监控的包或类名配置
            String monitorPackages = ConfigLoader.getMonitorPackages();
            String monitorClasses = ConfigLoader.getMonitorClasses();
            
            // 使用ByteBuddy的正确API：创建AgentBuilder
            AgentBuilder agentBuilder = new AgentBuilder.Default()
                    // 排除系统类和监控自身类
                    .ignore(ElementMatchers.nameStartsWith("com.jvmmonitor."))
                    .ignore(ElementMatchers.nameStartsWith("java."))
                    .ignore(ElementMatchers.nameStartsWith("javax."))
                    .ignore(ElementMatchers.nameStartsWith("sun."))
                    .ignore(ElementMatchers.nameStartsWith("com.sun."))
                    .ignore(ElementMatchers.nameStartsWith("net.bytebuddy."))
                    .ignore(ElementMatchers.nameStartsWith("org.slf4j."))
                    .ignore(ElementMatchers.nameStartsWith("ch.qos.logback."));
            
            // 默认不监控任何类，避免对主服务产生性能影响
            ElementMatcher.Junction<TypeDescription> typeMatcher = ElementMatchers.<TypeDescription>none();
            
            // 根据配置设置监控范围
            if (monitorPackages != null && !monitorPackages.isEmpty()) {
                // 监控指定的包
                logger.info("Monitoring specified packages: {}", monitorPackages);
                String[] packages = monitorPackages.split(",");
                
                ElementMatcher.Junction<TypeDescription> packageMatcher = null;
                for (String pkg : packages) {
                    String trimmedPkg = pkg.trim();
                    if (!trimmedPkg.isEmpty()) {
                        ElementMatcher.Junction<TypeDescription> currentMatcher = ElementMatchers.nameStartsWith(trimmedPkg);
                        if (packageMatcher == null) {
                            packageMatcher = currentMatcher;
                        } else {
                            packageMatcher = packageMatcher.or(currentMatcher);
                        }
                    }
                }
                
                if (packageMatcher != null) {
                    typeMatcher = packageMatcher;
                }
            } else if (monitorClasses != null && !monitorClasses.isEmpty()) {
                // 监控指定的类
                logger.info("Monitoring specified classes: {}", monitorClasses);
                String[] classes = monitorClasses.split(",");
                
                ElementMatcher.Junction<TypeDescription> classMatcher = null;
                for (String cls : classes) {
                    String trimmedCls = cls.trim();
                    if (!trimmedCls.isEmpty()) {
                        ElementMatcher.Junction<TypeDescription> currentMatcher = ElementMatchers.named(trimmedCls);
                        if (classMatcher == null) {
                            classMatcher = currentMatcher;
                        } else {
                            classMatcher = classMatcher.or(currentMatcher);
                        }
                    }
                }
                
                if (classMatcher != null) {
                    typeMatcher = classMatcher;
                }
            } else {
                logger.info("No monitoring packages or classes specified, defaulting to no class monitoring");
            }
            
            // 应用转换规则
            agentBuilder = agentBuilder
                .type(typeMatcher)
                .transform((builder, typeDescription, classLoader, module, protectionDomain) ->
                    builder
                        .method(ElementMatchers.isConstructor())
                        .intercept(MethodDelegation.to(ObjectSizeInterceptor.class))
                );
            
            // 确保agentBuilder已经有了transform配置
            // 我们已经在前面的条件分支中应用了transform，这里不再需要重复应用
            
            // 安装Agent到Instrumentation
            agentBuilder.installOn(inst);
            
            logger.info("Bytecode enhancement setup completed");
        } catch (Exception e) {
            logger.error("Failed to setup bytecode enhancement: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 启动定期统计任务
     */
    private void startScheduledTasks() {
        long interval = ConfigLoader.getMonitorInterval();
        
        scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread thread = new Thread(r, "jvm-monitor-scheduler");
            thread.setDaemon(true);
            return thread;
        });
        
        // 定期执行统计任务
        // 第一个参数为初始延迟，第二个参数为执行间隔
        // 使用0作为初始延迟，立即开始第一次执行
        scheduler.scheduleAtFixedRate(() -> {
            try {
                // 执行内存统计
                performMemoryStatistics();
            } catch (Exception e) {
                logger.error("Error in scheduled task: {}", e.getMessage(), e);
            }
        }, 0, interval, TimeUnit.MILLISECONDS);
        
        logger.info("Scheduled tasks started with interval: {}ms", interval);
    }
    
    /**
     * 执行内存统计
     */
    private void performMemoryStatistics() {
        boolean productionMode = ConfigLoader.isProductionMode();
        
        // 生产环境下使用更高级别的日志
        if (!productionMode) {
            logger.debug("Performing memory statistics...");
        }
        
        try {
            // 获取JMX内存统计 - 这是基础功能，所有模式下都执行
            jmxMonitor.collectMemoryStats();
            
            // 在非生产环境下执行完整的统计分析
            if (!productionMode) {
                // 获取对象大小统计
                objectSizeMonitor.collectStats();
                
                // 分析内存增长
                if (ConfigLoader.isMemoryGrowthAnalysisEnabled()) {
                    memoryAnalyzer.analyzeMemoryGrowth();
                }
                
                // 生成JVM优化建议
                if (ConfigLoader.isJvmAdviceEnabled()) {
                    memoryAnalyzer.generateJvmAdvice();
                }
            }
            
            // 重置统计计数器 - 无论什么模式都需要重置
            objectSizeMonitor.resetStats();
        } catch (Exception e) {
            // 生产环境下使用更高级别的日志
            if (productionMode) {
                logger.warn("Failed to perform memory statistics: {}", e.getMessage());
            } else {
                logger.error("Failed to perform memory statistics: {}", e.getMessage(), e);
            }
        }
    }
    
    /**
     * 获取Agent实例（单例）
     */
    public static JvmMonitorAgent getInstance() {
        return instance;
    }
    
    /**
     * 关闭Agent
     */
    public void shutdown() {
        logger.info("Shutting down JVM Monitor Agent...");
        
        try {
            // 关闭定时执行器
            if (scheduler != null) {
                scheduler.shutdown();
                if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                    scheduler.shutdownNow();
                }
            }
            
            // 关闭JMX监控
            if (jmxMonitor != null) {
                jmxMonitor.stop();
            }
            
            // 关闭内存分析器
            if (memoryAnalyzer != null) {
                // 假设MemoryAnalyzer类有stop方法，如果没有则可能需要添加
                // memoryAnalyzer.stop();
            }
            
            // 关闭对象大小监控
            if (objectSizeMonitor != null) {
                // 假设ObjectSizeMonitor类有stop方法，如果没有则可能需要添加
                // objectSizeMonitor.stop();
            }
            
            logger.info("JVM Monitor Agent shutdown completed");
        } catch (Exception e) {
            logger.error("Error during Agent shutdown: {}", e.getMessage(), e);
        }
    }
}